import os
import re
import nltk
from nltk.stem import PorterStemmer
from sklearn.feature_extraction.text import CountVectorizer
from collections import defaultdict

path = '../data'
files = os.listdir(path)


def get_text():
    texts = []
    for file in files:
        with open(path + '/' + file, encoding='utf-8') as f:
            texts.append(f.read())
    return texts


def get_vec(texts):
    vec = CountVectorizer()
    count = vec.fit_transform(texts)
    # print(vec.get_feature_names_out())
    return vec, count


def get_vec_pro(texts):
    stop_words = [
        'a', 'an', 'the', 'and', 'or', 'but', 'is', 'are', 'was', 'were', 'be', 'being', 'been',
        'have', 'has', 'had', 'do', 'does', 'did', 'will', 'would', 'shall', 'should', 'can', 'could',
        'may', 'might', 'must', 'of', 'at', 'by', 'for', 'with', 'in', 'on', 'to', 'from', 'as', 'into',
        'through', 'between', 'under', 'above', 'below', 'over', 'after', 'before', 'during', 'out', 'up',
        'down', 'off', 'over', 'again', 'then', 'once', 'here', 'there', 'now', 'how', 'all', 'any', 'no',
        'not', 'nor', 'only', 'so', 'very', 'too', 'just', 'more', 'most', 'less', 'least', 'many', 'few',
        'some', 'other', 'every', 'each', 'own', 'same', 'different', 'such', 'what', 'which', 'who', 'whom',
        'whose', 'where', 'when', 'why', 'how', 'until', 'while', 'after', 'before', 'because', 'if', 'unless',
        'since', 'while', 'whereas', 'although', 'though', 'even', 'once'
    ]
    vec = CountVectorizer(stop_words=stop_words)

    # 进行词干提取
    stemmer = PorterStemmer()
    texts_stemmed = [' '.join([stemmer.stem(word) for word in nltk.word_tokenize(text)]) for text in texts]
    count = vec.fit_transform(texts_stemmed)

    return vec, count


def gen_index(texts, vec, X):
    inverse_table = defaultdict(list)
    words = vec.get_feature_names_out()
    for paper_id, paper_value in enumerate(texts):
        papaer_offset = paper_value.strip().split('\n')[3].split(':')[1].strip()
        for word_id, word_value in enumerate(words):
            if X[paper_id][word_id] != 0:
                position = [p.span() for p in re.finditer(r'\b' + word_value + r'\b', papaer_offset)]
                inverse_table[word_value].append((paper_id, X[paper_id][word_id], position, word_value))
    # key : value
    # word : (paper_id, freq, (start, end), word_value)
    return inverse_table
